import torch
from torch import nn


# ----- 3x28x28  --> nx7x7 (6x6、8x8)
# ----- 定义一个简单的神经网络 ------
class CustomNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, 1, 1)
        self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
        # 池化
        self.pool = nn.MaxPool2d(2, 2)

    # 前向传播
    def forward(self, x):
        x = self.pool(self.conv1(x))
        x = self.pool(self.conv2(x))
        return x


# (1, 3, 28, 28) 1:代表这个图像有多少个（批次 batch_size）
# (1, 3, 28, 28) 3:代表这个图像有多少个通道（channel）
# (1, 3, 28, 28) 28，28:代表这个图像的大小（image_size）
image = torch.randn(1, 3, 28, 28)
# image = torch.randn(1, 3, 100, 100) 这是水果分类的图像大小
model = CustomNet()
result = model(image)
print(result.shape)  # 输出结果的形状
